import argparse
import os
from pytorch_grad_cam.guided_backprop import DeconvReLUModel, GothroughReLUModel
import cv2
import numpy as np
import torch
from torchvision import models
import pdb
import copy
import timm
from functools import partial

from pytorch_grad_cam import GradCAM, \
    ScoreCAM, \
    GradCAMPlusPlus, \
    AblationCAM, \
    XGradCAM, \
    EigenCAM, \
    EigenGradCAM

from pytorch_grad_cam import GuidedBackpropReLUModel
from pytorch_grad_cam.utils.image import show_cam_on_image, \
    deprocess_image, \
    preprocess_image


def get_args():
    parser = argparse.ArgumentParser()
    # parser.add_argument('--use-cuda', action='store_true', default=False,
    #                     help='Use NVIDIA GPU acceleration')
    parser.add_argument('--use-cuda', type=str, default='0',
                        help='Use NVIDIA GPU acceleration, gpu list')
    parser.add_argument('--image-path', type=str, default='./examples/both.png',
                        help='Input image path')
    parser.add_argument('--aug_smooth', action='store_true',
                        help='Apply test time augmentation to smooth the CAM')
    parser.add_argument('--eigen_smooth', action='store_true',
                        help='Reduce noise by taking the first principle componenet'
                        'of cam_weights*activations')
    parser.add_argument('--method', type=str, default='gradcam',
                        choices=['gradcam', 'gradcam++', 'scorecam', 'xgradcam',
                                 'ablationcam', 'eigencam', 'eigengradcam'],
                        help='Can be gradcam/gradcam++/scorecam/xgradcam'
                             '/ablationcam/eigencam/eigengradcam')
    parser.add_argument('--model_name', type=str, default='resnet50')
    args = parser.parse_args()

    return args


if __name__ == '__main__':
    """ python cam.py -image-path <path_to_image>
    Example usage of loading an image, and computing:
        1. CAM
        2. Guided Back Propagation
        3. Combining both
    """

    args = get_args()
    use_cuda = args.use_cuda != '-1' and torch.cuda.is_available()
    if use_cuda:
        os.environ["CUDA_VISIBLE_DEVICES"] = args.use_cuda
        print('Using GPU for acceleration')
    else:
        print('Using CPU for computation')

    methods = \
        {"gradcam": GradCAM,
         "scorecam": ScoreCAM,
         "gradcam++": GradCAMPlusPlus,
         "ablationcam": AblationCAM,
         "xgradcam": XGradCAM,
         "eigencam": EigenCAM,
         "eigengradcam": EigenGradCAM}

    adv_models = ['adv_inception_v3', 'tf_efficientnet_b0_ap', 'ens_adv_inception_resnet_v2']
    model_name = args.model_name
    if model_name in adv_models:
        model_ = partial(timm.create_model, model_name)
    else:
        model_ = getattr(models, model_name)
    model = model_(pretrained=True)
    # pdb.set_trace()
    model_layer_dict = {"resnet18": ["layer1", "layer2", "layer3", "layer4"],
                        "resnet50": ["layer1", "layer2", "layer3", "layer4"],
                        "vgg16":["features"],
                        "vgg19":["features"],
                        "densenet161":["features"],
                        "mnasnet1_0":["layers"],
                        "adv_inception_v3":["Mixed_7c"],
                        "ens_adv_inception_resnet_v2":["conv2d_7b"],
                        "tf_efficientnet_b0_ap":["blocks", "conv_head"],
                        }
    layers = model_layer_dict[model_name]
    # target_layer = model.layer4[-1]
    

    # Choose the target layer you want to compute the visualization for.
    # Usually this will be the last convolutional layer in the model.
    # Some common choices can be:
    # Resnet18 and 50: model.layer4[-1]
    # VGG, densenet161: model.features[-1]
    # mnasnet1_0: model.layers[-1]
    # You can print the model to help chose the layer

    rgb_img = cv2.imread(args.image_path, 1)[:, :, ::-1]
    rgb_img = np.float32(rgb_img) / 255
    input_tensor = preprocess_image(rgb_img, mean=[0.485, 0.456, 0.406],
                                    std=[0.229, 0.224, 0.225])
    # If None, returns the map for the highest scoring category.
    # Otherwise, targets the requested category.
    target_category = None

    gb_model = GuidedBackpropReLUModel(
        model=copy.deepcopy(model), use_cuda=use_cuda)
    # gb_model = DeconvReLUModel(model=model, use_cuda=use_cuda)
    # gb_model = GothroughReLUModel(model=copy.deepcopy(model), use_cuda=use_cuda)
    gb = gb_model(input_tensor, target_category=target_category)
    cv2.imwrite(f'outputs/gb_{args.method}_{model_name}.jpg', deprocess_image(gb))
    
    # pdb.set_trace()
    for layer in reversed(layers):
        target_module = getattr(model, layer)
        if not hasattr(target_module, '__iter__'):
            target_module = [target_module]
        for i, target_layer in enumerate(reversed(target_module)):
            str_name = "{}_last{}".format(layer, i + 1)
            cam = methods[args.method](model=model,
                                       target_layer=target_layer,
                                       use_cuda=use_cuda)

            # AblationCAM and ScoreCAM have batched implementations.
            # You can override the internal batch size for faster computation.
            cam.batch_size = 32

            grayscale_cam = cam(input_tensor=input_tensor,
                                target_category=target_category,
                                aug_smooth=args.aug_smooth,
                                eigen_smooth=args.eigen_smooth)

            # Here grayscale_cam has only one image in the batch
            grayscale_cam = grayscale_cam[0, :]

            cam_image = show_cam_on_image(rgb_img, grayscale_cam)

            cam_mask = cv2.merge([grayscale_cam, grayscale_cam, grayscale_cam])
            cam_gb = deprocess_image(cam_mask * gb)

            # cv2.imwrite(f'outputs/{args.method}_cam.jpg', cam_image)
            # cv2.imwrite(f'outputs/{args.method}_deconv.jpg', gb)
            # cv2.imwrite(f'outputs/{args.method}_cam_deconv.jpg', cam_gb)
            cv2.imwrite(f'outputs/cam_{args.method}_{model_name}_{str_name}.jpg', cam_image)
            cv2.imwrite(f'outputs/cam_gb_{args.method}_{model_name}_{str_name}.jpg', cam_gb)

            del cam
            del grayscale_cam
            del cam_image
            del cam_mask
            torch.cuda.empty_cache()
            break
        break